# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
import random
import numpy as np

from typing import Union, List
from atk.case_generator.generator.generate_types import GENERATOR_REGISTRY
from atk.case_generator.generator.base_generator import CaseGenerator
from atk.configs.case_config import InputCaseConfig, CaseConfig


@GENERATOR_REGISTRY.register("ascend_generate_aclnn_moe_token_permute_with_ep")
class DefaultGenerator(CaseGenerator):
    def __init__(self, config):
        super().__init__(config)
        self.tokens_num = 1
        self.hidden_size = 1
        self.topk = 1
        self.out_token_nums = 1
        self.start = 0
        self.end = 0

    def after_input_config(
            self,
            index: int,
            input_case: Union[InputCaseConfig, List[InputCaseConfig]]
    ) -> Union[InputCaseConfig, List[InputCaseConfig]]:

        if index == 0: # tokens
            self.tokens_num = input_case.shape[0]
            self.hidden_size = input_case.shape[1]

        if index == 1: # indices
            if len(input_case.shape) == 2:
                self.topk = input_case.shape[1]
                if self.topk > 512:
                    self.topk = 512
            self.out_token_nums = self.tokens_num * self.topk

        if index == 3: # rangeOptional
            self.start = random.randint(0, self.out_token_nums)
            self.end = random.randint(0, self.out_token_nums)
            if self.end < self.start:
                self.start, self.end = self.end, self.start
            if self.start == self.end:
                self.start = 0
                self.end = self.out_token_nums
            input_case[0].range_values = self.start
            input_case[1].range_values = self.end

        return input_case

    def after_case_config(self, case_config: CaseConfig) -> CaseConfig:
        case_config.inputs[0].shape = [self.tokens_num, self.hidden_size]
        case_config.inputs[1].shape = [self.tokens_num, self.topk]
        case_config.inputs[2].shape = [self.tokens_num, self.topk]
        case_config.inputs[2].dtype = case_config.inputs[0].dtype
        return case_config